package main

import (
	"bytes"
	"encoding/binary"
	. "gotunnel/tunnel"
	"log"
	"net"
	"sync"
	"time"
)

var (
	remote_bind             = "192.168.128.1:8090"
	local_uppon             = "127.0.0.1:2012"
	remote_bind_client      net.Conn
	local_uppon_clients_map map[uint16]net.Conn = make(map[uint16]net.Conn, 100)
	buffer                                      = sync.Pool{
		New: func() interface{} { return make([]byte, BUFF_SIZE) },
	}
	buffer_new = sync.Pool{
		New: func() interface{} { return new(bytes.Buffer) },
	}
	pool_conn = sync.Pool{
		New: func() interface{} { return new(net.TCPConn) },
	}
)

func getUpponConnBySessionId(session_id uint16) net.Conn {
	for k, v := range local_uppon_clients_map {
		if k == session_id {
			return v
		}
	}
	return nil
}
func getSessionIdByUpponConn(c net.Conn) uint16 {
	for k, v := range local_uppon_clients_map {
		if c != nil && v != nil && c == v {
			return k
		}
	}
	return 0
}
func clearUpponClients() {
	for k, v := range local_uppon_clients_map {
		if v != nil {
			v.Close()
			delete(local_uppon_clients_map, k)
			log.Println("远程映射服务器关闭，断开映射", k)
		}
	}
}
func pack(conn net.Conn, token uint8, session_id uint16, data []byte) {
	var buf = buffer_new.Get().(*bytes.Buffer)
	buf.WriteByte(token)
	binary.Write(buf, binary.LittleEndian, session_id)
	buf.Write(data)
	if conn != nil {
		buf.WriteTo(conn)
	}
}
func OnTunnelLocalUpponClient(c net.Conn, session_id uint16) {
	defer c.Close()
	log.Println("更新映射", session_id)
	local_uppon_clients_map[session_id] = c
	for {
		data := buffer.Get().([]byte)
		n, err := c.Read(data)
		if err != nil {
			log.Println("本地映射断开", session_id)
			log.Println("删除映射", session_id)
			pack(remote_bind_client, TOKEN_SERVER_DISCONNECT, session_id, []byte{})
			delete(local_uppon_clients_map, session_id)
			break
		}
		if remote_bind_client != nil {
			log.Println("开始回复转发数据", session_id, n)
			pack(remote_bind_client, TOKEN_SERVER_REFORWARD, session_id, data[:n])
		} else {
			log.Println("重转发失败,可能远程服务器已关闭")
		}
	}
}
func OnTunnelRemoteClient(c net.Conn, disconnect chan bool) {
	defer c.Close()
	c.(*net.TCPConn).SetKeepAlive(true)
	remote_bind_client = c

	for {
		data := buffer.Get().([]byte)
		n, err := c.Read(data)
		if err != nil {
			remote_bind_client.Close()
			remote_bind_client = nil
			disconnect <- true
			break
		}
		session_id := binary.LittleEndian.Uint16(data[1:3])
		var conn = getUpponConnBySessionId(session_id)
		switch data[0] {
		case TOKEN_SERVER_DISCONNECT:
			log.Println("远程映射服务断开连接", session_id)
			conn.Close()
			delete(local_uppon_clients_map, session_id)
		case TOKEN_SERVER_FORWARD:
			log.Println("收到服务端转发数据", session_id, n)
			if conn == nil {
				if session_id > 0 {
					log.Println("连接本地映射服务端", local_uppon, session_id)
					var client net.Conn = pool_conn.Get().(*net.TCPConn)
					var err error
					client.(*net.TCPConn).SetKeepAlive(true)
					client, err = net.Dial("tcp", local_uppon)
					if err != nil {
						log.Println("连接本地映射服务端失败", local_uppon, session_id)
						pack(remote_bind_client, TOKEN_SERVER_DISCONNECT, session_id, []byte{})
						continue
					}
					client.Write(data[3:n])
					go OnTunnelLocalUpponClient(client, session_id)
				}

			} else {
				conn.Write(data[3:n])
			}
		default:
		}
	}
}
func main() {
	var quit = make(chan bool)
	go func() {
		for {
			log.Println("开始连接远程映射服务器", remote_bind)
			conn, err := net.Dial("tcp", remote_bind)
			if err != nil {
				log.Println("连接远程映射服务器失败", remote_bind)
				time.Sleep(time.Millisecond * 1200)
				continue
			}
			log.Println("连接远程映射服务器成功", remote_bind)
			pack(conn, TOKEN_SERVER_DISCONNECT_ALL, 0, []byte{})
			var disconnect = make(chan bool)
			go OnTunnelRemoteClient(conn, disconnect)
			<-disconnect

		}
	}()
	go func() {
		for {
			select {
			case <-time.After(time.Second * 5):
				if remote_bind_client == nil {
					clearUpponClients()
				}
			}
		}
	}()
	go func() {
		for {
			select {

			case <-time.After(time.Millisecond * 200):
				for k, _ := range local_uppon_clients_map {
					if remote_bind_client != nil {
						pack(remote_bind_client, TOKEN_SERVER_STATUS, k, []byte{})
					}

				}

			}
		}
	}()

	<-quit
}
